#include "fsw.h"

static fftw_plan	fplan, bplan, fpp, fpz, bpp, bpz, bpr, fpr;
static fftw_complex	p2, r2, *mul, *dm, *rm, *tm, *dbg; 
static int		i, j, k, N, l, h, kcut, k_pow;
static double		nzu, g, av, dt, dx, eps, depth, ynot, depth2;
static double		re, im, depth_old;
static arr2_str		r[4], zs, s;

void derivative(ctrl_ptr in)
{ 
  fftw_execute(fpp);
  fftw_execute(fpz);
  for (j = 0; j < N; j++)
  {
    if (j > N/2) rm[j]  = 1.I*tanh((2*(in->h_curr)*pi/l)*(j - N));
    else         rm[j]  = 1.I*tanh((2*(in->h_curr)*pi/l)*j);
    zs.p[j] = dm[j]*zs.p[j];
    zs.z[j] = dm[j]*zs.z[j];
    zs.r[j] = rm[j]*zs.p[j];
  }
  fftw_execute(bpp);
  fftw_execute(bpz);
  fftw_execute(bpr);
}
void t_transform(ctrl_ptr in)
{ 
  fftw_execute(fpp);
  for (j = 0; j < N; j++) {
    if (j > N/2) tm[j]  = -1.I/tanh((2*(in->h_curr)*pi/l)*(j - N))/N;
    else 	 tm[j]  = -1.I/tanh((2*(in->h_curr)*pi/l)*j)/N;
    zs.p[j] = tm[j]*zs.p[j]; 
  }
  zs.p[0] = 0;
  fftw_execute(bpp);
}
void r_transform(ctrl_ptr in)
{ 
  fftw_execute(fpp);
  for (j = 0; j < N; j++) {
    if (j > N/2) rm[j]  = 1.I*tanh((2*(in->h_curr)*pi/l)*(j - N));
    else         rm[j]  = 1.I*tanh((2*(in->h_curr)*pi/l)*j);
    zs.p[j] = rm[j]*zs.p[j]/N; 
  }
  fftw_execute(bpp);
}

//--------------------------------------------------------------------
void init_fft(ctrl_ptr ctrl)
{
  N     = ctrl->n; l  = ctrl->l; h = ctrl->h; kcut = ctrl->kcut; dx = ctrl->dx;
  g	= ctrl->g; av = (ctrl->av); dt = ctrl->dt; k_pow = ctrl->k_power; eps = ctrl->eps;
  for (k = 0; k < 4; k++ ) {
    r[k].p = malloc(sizeof(fftw_complex)*(ctrl->n));
    r[k].z = malloc(sizeof(fftw_complex)*(ctrl->n));
  }
  zs.p = malloc(sizeof(fftw_complex)*(ctrl->n));
  zs.z = malloc(sizeof(fftw_complex)*(ctrl->n));
  zs.r = malloc(sizeof(fftw_complex)*(ctrl->n));
  s.p  = malloc(sizeof(fftw_complex)*(ctrl->n));
  s.z  = malloc(sizeof(fftw_complex)*(ctrl->n));
  mul  = malloc(sizeof(fftw_complex)*(ctrl->n));
  dm   = malloc(sizeof(fftw_complex)*(ctrl->n));
  tm   = malloc(sizeof(fftw_complex)*(ctrl->n));
  rm   = malloc(sizeof(fftw_complex)*(ctrl->n));
  dbg  = malloc(sizeof(fftw_complex)*(ctrl->n));
  for ( j = 0; j < N; j++ ){
    if (j > N/2) {
	p2 = pow(2*pi*(j-N)/l,k_pow);
	dm[j]  = 1.I*(2*pi/l)*(j - N)/N;
    } else {
        p2 = pow(2*pi*j/l,k_pow);
	dm[j] = 1.I*(2*pi/l)*j/N;
    }
    mul[j]  = exp(-av*p2*dt)/N;
  }
  mul[0] = 0;  if (eps == 0) { rm[0] = 0; tm[0]  = 0; } depth = h; depth2 = 0;
  depth_old = 0;
  fpp   = fftw_plan_dft_1d(ctrl->n, zs.p, zs.p, FFTW_FORWARD, fmode);
  fpz   = fftw_plan_dft_1d(ctrl->n, zs.z, zs.z, FFTW_FORWARD, fmode);
  fpr   = fftw_plan_dft_1d(ctrl->n, zs.r, zs.r, FFTW_FORWARD, fmode);
  bpp   = fftw_plan_dft_1d(ctrl->n, zs.p, zs.p, FFTW_BACKWARD, fmode);
  bpz   = fftw_plan_dft_1d(ctrl->n, zs.z, zs.z, FFTW_BACKWARD, fmode);
  bpr   = fftw_plan_dft_1d(ctrl->n, zs.r, zs.r, FFTW_BACKWARD, fmode);
}
//-------------------------------------------------------------------
void compute_billet(ctrl_ptr in)
{
  for (j = 0; j < N; j++) {
    zs.p[j] = in->psi[j];
    zs.z[j] = in->z[j];
  }
  derivative(in);
  for (j = 0; j < N; j++) {
    in->dps[j] = creal(zs.p[j]);
    in->dz[j]  = 1 + zs.z[j];
    in->rps[j] = creal(zs.r[j]);
    nzu = cabs(in->dz[j])*cabs(in->dz[j]);  
    zs.p[j] = (in->rps[j])/nzu;
  }
  t_transform(in);  
  for (j = 0; j < N; j++) {
    in->trpj[j] = creal(zs.p[j]);
    zs.p[j] = in->psi[j];
    zs.z[j] = in->z[j];
    //printf("%e\t%eI\n",creal(in->trpj[j]),cimag(in->trpj[j]));
  }
  //exit(1);
  fftw_execute(fpz);
  fftw_execute(fpp);
  for (j = 0; j < N; j++) {
    in->fps[j] = zs.p[j]/N;
    in->fz[j] = zs.z[j]/N;
  }
}

void rhs_one(ctrl_ptr in, int num)
{
 put_zero_mode(in);
 compute_billet(in); 
 for ( j = 0; j < N; j++)
 {
   nzu = cabs(in->dz[j])*cabs(in->dz[j]);
   p2 = creal((in->dps[j])*(in->dps[j]));
   r2 = creal((in->rps[j])*(in->rps[j]));
   r[num].p[j] = -g*cimag(in->z[j]) - (0.5/nzu)*(p2-r2) - (in->dps[j])*(in->trpj[j]);
   //r[num].p[j] = -g*(cimag(in->z[j])-((in->h_curr)-h)) - (0.5/nzu)*(p2-r2) - creal((in->dps[j])*(in->trpj[j]));
   r[num].z[j] = -(in->dz[j])*((in->trpj[j]) + 1.I*(in->rps[j])/nzu);
 }  
}
void rhs_two(ctrl_ptr in, int num)
{
  for ( j = 0; j < N; j ++) {
    zs.p[j] = in->psi[j];
  
  }
}

/*void put_zero_mode(ctrl_ptr in)  
{
  for ( j = 0; j < N; j++) {
    zs.z[j] = in->z[j];
    zs.p[j] = in->dz[j];
    zs.r[j] = in->psi[j];
  }
  fftw_execute(fpz); 
  fftw_execute(fpp);
  fftw_execute(fpr);
  depth = 0; 
  for (j = 0; j < N; j++) {
    zs.z[j] = zs.z[j]/N;
    zs.p[j] = zs.p[j]/N;
    zs.r[j] = zs.r[j]/N;
  }
  for (j = 1; j < N; j++) depth = depth - N*(0.5I*dm[j])*pow(cabs(zs.z[j]),2); 
  depth = depth + h;
  //zs.z[0] = 0 + 1.I*(depth-h); 

  in->mean_level = 0; in->k_energy = 0; 
  for (j = 0; j < N; j++) {
    if (j > N/2) rm[j]  = (eps + 1.I*tanh((2*(in->h_curr)*pi/l)*(j - N)));
    else         rm[j]  = (eps + 1.I*tanh((2*(in->h_curr)*pi/l)*j));
    in->mean_level += zs.p[j]*conj(zs.z[j]); 
    in->k_energy += -0.5*zs.r[j]*conj(zs.r[j]*dm[j]*rm[j]*N);
  }
  //in->mean_level = in->mean_level - ()
  fftw_execute(bpz);
  for (j = 0; j < N; j++) in->z[j] = zs.z[j];
  in->h_curr = depth;
}*/

void rk4(ctrl_ptr in, double multipl)
{
  rhs_one(in,0);
  for (j = 0; j < N; j++){
    s.p[j] = in->psi[j];
    s.z[j] = in->z[j];
    in->psi[j] = s.p[j] + 0.5*dt*creal(r[0].p[j]);
    in->z[j] = s.z[j] + 0.5*dt*(r[0].z[j]);
  }
  rhs_one(in,1);
  for (j = 0; j < N; j++){
    in->psi[j] = s.p[j] + 0.5*dt*(r[1].p[j]);
    in->z[j] = s.z[j] + 0.5*dt*(r[1].z[j]);
  }
  rhs_one(in,2);
  for (j = 0; j < N; j++){
    in->psi[j] = s.p[j] + dt*(r[2].p[j]);
    in->z[j] = s.z[j] + dt*(r[2].z[j]);
  }
  rhs_one(in,3);
  for (j = 0; j < N; j++){
    in->psi[j] = s.p[j] + multipl*(dt/6.)*(r[0].p[j] + 2.*r[1].p[j] + 2.*r[2].p[j] + r[3].p[j]);
    in->z[j]   = s.z[j] + multipl*(dt/6.)*(r[0].z[j] + 2.*r[1].z[j] + 2.*r[2].z[j] + r[3].z[j]);
    zs.p[j] = in->psi[j];
    zs.z[j] = in->z[j];
  }
  fftw_execute(fpp);
  fftw_execute(fpz);
  for (j = 0; j < N; j++)
  {
    zs.p[j] = mul[j]*zs.p[j];
    zs.z[j] = mul[j]*zs.z[j];
  }
  fftw_execute(bpp);
  fftw_execute(bpz);
  in->p_energy = 0; 
  for (j = 0; j < N; j++) {
    in->psi[j] = creal(zs.p[j]);
    in->z[j]   = zs.z[j];
  }
  //put_zero_mode(in);
  in->k_energy_x = 0;
  for (j = 0; j < N; j++) {
    p2 = cimag(in->z[j])*(cimag(in->z[j]))*creal(in->dz[j]);
    //p2 = (cimag(in->z[j])-(depth-h))*(cimag(in->z[j])-(depth-h))*creal(in->dz[j]);
    r2 = cabs((in->psi[j])*(in->rps[j]));
    in->k_energy_x = in->k_energy + 0.5*r2*dx;
    in->p_energy = in->p_energy + 0.5*g*p2*dx/l;
    //in->mean_level = in->mean_level + creal(in->dz[j])*cimag(in->z[j])*dx;
  }
  in->t_energy = in->k_energy + in->p_energy;
  in->h_curr = depth;
}




